import acl
import numpy as np
from aclnet.constant import PIXEL_FORMAT_YUV_SEMIPLANAR_420, ACL_ERROR_NONE, ACL_MEMCPY_HOST_TO_DEVICE
from aclnet.constant import NMat, check_ret

class Dvpp(object):
    def __init__(self, context):
        self.context = context

        self.stream, ret = acl.rt.create_stream()
        check_ret("[dvpp]acl.rt.create_stream", ret)

        self.dvpp_channel_desc = acl.media.dvpp_create_channel_desc()
        acl.media.dvpp_create_channel(self.dvpp_channel_desc)

    def __del__(self):
        if self.dvpp_channel_desc:
            acl.media.dvpp_destroy_channel(self.dvpp_channel_desc)
            acl.media.dvpp_destroy_channel_desc(self.dvpp_channel_desc)
            self.dvpp_channel_desc = None

    def imread(self, filepath):
        data = np.fromfile(filepath, dtype=np.uint8)
        ptr = acl.util.numpy_to_ptr(data)
        width, height, components, ret = acl.media.dvpp_jpeg_get_image_info(ptr, data.size)
        check_ret("[dvpp]acl.media.dvpp_jpeg_get_image_info", ret)

        decode_out_width = ((width + 127) // 128) * 128
        decode_out_height = ((height + 15) // 16) * 16
        decode_out_buffer_size = decode_out_width * decode_out_height * 3 // 2
        buffer, ret = acl.media.dvpp_malloc(decode_out_buffer_size)
        check_ret("acl.media.dvpp_malloc", ret)

        input_buffer, ret = acl.media.dvpp_malloc(data.size)
        check_ret("acl.media.dvpp_malloc", ret)
        acl.rt.memcpy(input_buffer, data.size, ptr, data.size, ACL_MEMCPY_HOST_TO_DEVICE)

        output_dvpp_pic_desc = acl.media.dvpp_create_pic_desc()
        acl.media.dvpp_set_pic_desc_data(output_dvpp_pic_desc, buffer)
        acl.media.dvpp_set_pic_desc_format(output_dvpp_pic_desc, PIXEL_FORMAT_YUV_SEMIPLANAR_420)
        acl.media.dvpp_set_pic_desc_width(output_dvpp_pic_desc, width)
        acl.media.dvpp_set_pic_desc_height(output_dvpp_pic_desc, height)
        acl.media.dvpp_set_pic_desc_width_stride(output_dvpp_pic_desc, decode_out_width)
        acl.media.dvpp_set_pic_desc_height_stride(output_dvpp_pic_desc, decode_out_height)
        acl.media.dvpp_set_pic_desc_size(output_dvpp_pic_desc, decode_out_buffer_size)
        ret = acl.media.dvpp_jpeg_decode_async(self.dvpp_channel_desc, input_buffer, data.size, output_dvpp_pic_desc, self.stream)
        if ret != ACL_ERROR_NONE:
            raise Exception("dvpp_jpeg_decode_async failed ret={}".format(ret))
        acl.rt.synchronize_stream(self.stream)
        return NMat(buffer, decode_out_buffer_size, width, height, decode_out_width, decode_out_height)

    def resize(self, img, width, height):
        resize_input_desc = acl.media.dvpp_create_pic_desc()
        acl.media.dvpp_set_pic_desc_data(resize_input_desc, img.ptr_dev)
        acl.media.dvpp_set_pic_desc_format(resize_input_desc, PIXEL_FORMAT_YUV_SEMIPLANAR_420)
        acl.media.dvpp_set_pic_desc_width(resize_input_desc, img.width)
        acl.media.dvpp_set_pic_desc_height(resize_input_desc, img.height)
        acl.media.dvpp_set_pic_desc_width_stride(resize_input_desc, img.stride_width)
        acl.media.dvpp_set_pic_desc_height_stride(resize_input_desc, img.stride_height)
        acl.media.dvpp_set_pic_desc_size(resize_input_desc, img.data_size)

        resize_stride_width = ((width + 127) // 128) * 128
        resize_stride_height = ((height + 15) // 16) * 16
        resize_data_size = resize_stride_width * resize_stride_height * 3 // 2
        resize_buffer, ret = acl.media.dvpp_malloc(resize_data_size)
        check_ret("acl.media.dvpp_malloc", ret)

        resize_output_desc = acl.media.dvpp_create_pic_desc()
        acl.media.dvpp_set_pic_desc_data(resize_output_desc, resize_buffer)
        acl.media.dvpp_set_pic_desc_format(resize_output_desc, PIXEL_FORMAT_YUV_SEMIPLANAR_420)
        acl.media.dvpp_set_pic_desc_width(resize_output_desc, width)
        acl.media.dvpp_set_pic_desc_height(resize_output_desc, height)
        acl.media.dvpp_set_pic_desc_width_stride(resize_output_desc, resize_stride_width)
        acl.media.dvpp_set_pic_desc_height_stride(resize_output_desc, resize_stride_height)
        acl.media.dvpp_set_pic_desc_size(resize_output_desc, resize_data_size)

        crop_area = acl.media.dvpp_create_roi_config(0,
                                                     img.width if img.width % 2 == 1 else img.width - 1,
                                                     0,
                                                     img.height if img.height % 2 == 1 else img.height - 1)  # left,right,top,bottom

        paste_area = acl.media.dvpp_create_roi_config(0,
                                                      width if width % 2 == 1 else width - 1,
                                                      0,
                                                      height if height % 2 == 1 else height - 1)  # left,right,top,bottom

        ret = acl.media.dvpp_vpc_crop_and_paste_async(self.dvpp_channel_desc,
                                              resize_input_desc,
                                              resize_output_desc,
                                              crop_area,
                                              paste_area,
                                              self.stream)
        acl.rt.synchronize_stream(self.stream)
        check_ret("[dvpp]acl.media.dvpp_vpc_resize_async", ret)
        return NMat(resize_buffer, resize_data_size, width, height, resize_stride_width, resize_stride_height)

# dvpp = Dvpp(0, False)
# image = dvpp.imread("./data/moto.jpg")
# image = dvpp.resize(image, 2048, 1536)
# import cv2
# img = cv2.cvtColor(image.numpy(),cv2.COLOR_YUV2BGR_NV12)
# cv2.imwrite('1.jpg', img)